#include "unionfind.h"

#include <stdlib.h>
#include<stdio.h>
#include<assert.h>
const uint32_t SZ = 2;

uint32_t get_set_size(union_find_t * uf, const uint32_t idx)
{
    return uf->data[SZ*uf->get_rep(uf, idx) +1];
}

uint32_t get_rep(union_find_t * uf,  const uint32_t idx)
{
    // Terminal Case: a node is its own parent
    if (uf->data[SZ*idx] == idx)
        return idx;

    uint32_t root = uf->get_rep(uf, uf->data[SZ*idx]);

    uf->data[SZ*idx] = root; // cache the result

    return root;
}

uint32_t connect(union_find_t * uf,  const uint32_t ida, const uint32_t idb)
{

    uint32_t aroot = uf->get_rep(uf, ida);
    uint32_t broot = uf->get_rep(uf, idb);

    if (aroot == broot)
        return aroot;

    int asz = uf->data[SZ*aroot + 1];
    int bsz = uf->data[SZ*broot + 1];

    if (asz > bsz) {
        uf->data[SZ*broot] = aroot;
        uf->data[SZ*aroot + 1] += bsz;
        return aroot;
    } else {
        uf->data[SZ*aroot] = broot;
        uf->data[SZ*broot + 1] += asz;
        return broot;
    }
}


union_find_t * union_find_init(const uint32_t sz)
{
    union_find_t * uf = malloc(sizeof(union_find_t));
    uf->sz = sz;
    uf->data = calloc(sizeof(uint32_t), uf->sz*SZ);

    for (uint32_t i = 0; i < uf->sz; i++) {
        uf->data[SZ*i + 0] = i;
        uf->data[SZ*i + 1] = 1;
    }

    uf->get_set_size = get_set_size;
    uf->get_rep = get_rep;
    uf->connect = connect;

    return uf;
}

void union_find_destroy(union_find_t * uf)
{
    free(uf->data);
    free(uf);
}
